miχPODS: New baseline#

TODO

  • more intake usage

  • why does the stability diagram load multiple times

  • hvplot defaults

    • turn off scroll zoom

    • copy over presentation stuff

Catalog display#

Hide code cell source
%load_ext rich

# MOM6 run catalog
catalog = {
    "baseline": (
        "baseline",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods",
    ),
    "epbl": ("ePBL", "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.epbl.001.mixpods"),
    "kpp.lmd.002": (
        "KPP Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.002.mixpods",
    ),
    "kpp.lmd.003": (
        "KPP Ri0=0.5, Ric=0.2,",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.003.mixpods",
    ),
    "kpp.lmd.004": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods",
    ),
    "baseline.N150": (
        "baseline N=150",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N150.baseline.mixpods",
    ),
    "kpp.lmd.004.N150": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5, N=150",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N150.kpp.lmd.004.mixpods",
    ),
    "new_baseline.hb": (
        "KD=0, KV=0",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb",
    ),
    "new_baseline.kpp.lmd.004": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods",
    ),
    "new_baseline.kpp.lmd.005": (
        "KPP ν0=2.5, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods",
    ),
}
catalog
The rich extension is already loaded. To reload it, use:
  %reload_ext rich
{
    'baseline': ('baseline', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods'),
    'epbl': ('ePBL', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.epbl.001.mixpods'),
    'kpp.lmd.002': (
        'KPP Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.002.mixpods'
    ),
    'kpp.lmd.003': (
        'KPP Ri0=0.5, Ric=0.2,',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.003.mixpods'
    ),
    'kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods'
    ),
    'baseline.N150': ('baseline N=150', 'gmom.e23.GJRAv3.TL319_t061_zstar_N150.baseline.mixpods'),
    'kpp.lmd.004.N150': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5, N=150',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N150.kpp.lmd.004.mixpods'
    ),
    'new_baseline.hb': ('KD=0, KV=0', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb'),
    'new_baseline.kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.kpp.lmd.005': (
        'KPP ν0=2.5, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods'
    )
}

Setup#

Hide code cell source
%load_ext autoreload
%load_ext rich
%load_ext watermark

import warnings

import cf_xarray as cfxr
import distributed
import holoviews as hv
import matplotlib.pyplot as plt
import math
import numpy as np
import pandas as pd
import tqdm
from datatree import DataTree

import xarray as xr

%aimport pump
from pump import mixpods

hv.notebook_extension("bokeh")

cfxr.set_options(warn_on_missing_variables=False)
xr.set_options(keep_attrs=True, display_expand_data=False)

plt.style.use("bmh")
plt.rcParams["figure.dpi"] = 140


%watermark -iv
Hide code cell output
The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/dask_jobqueue/core.py:20: FutureWarning: tmpfile is deprecated and will be removed in a future release. Please use dask.utils.tmpfile instead.
  from distributed.utils import tmpfile
pandas       : 1.5.3
sys          : 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:08:06) [GCC 11.3.0]
xarray       : 2023.3.0
json         : 2.0.9
tqdm         : 4.65.0
cf_xarray    : 0.8.0
pump         : 1.0+247.g1f1c5e1.dirty
numpy        : 1.23.5
distributed  : 2023.3.0
dask_jobqueue: 0.7.3
matplotlib   : 3.7.1
holoviews    : 1.15.4
Hide code cell source
if "client" in locals():
    client.close()
    del client
if "cluster" in locals():
    cluster.close()

import ncar_jobqueue

cluster = ncar_jobqueue.NCARCluster(
    local_directory="/local_scratch/pbs.$PBS_JOBID/dask/spill",
)
cluster.adapt(minimum_jobs=1, maximum_jobs=4)
client = distributed.Client(cluster)
client
Hide code cell output

Client

Client-6366c295-d427-11ed-8f74-3cecef1b11fa

Connection method: Cluster object Cluster type: dask_jobqueue.PBSCluster
Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/dcherian/proxy/8787/status

Cluster Info

Read#

Always read TAO#

tao_gridded = mixpods.load_tao()
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "depth" starting at index 42. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "time" starting at index 199726. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "longitude" starting at index 2. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(

Subset Catalog, build tree#

catalog_sub = {
    k: catalog[k]
    for k in catalog.keys()
    if k == "kpp.lmd.004" or ("baseline" in k and "150" not in k)
}
display(catalog_sub)
{
    'baseline': ('baseline', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods'),
    'kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.hb': ('KD=0, KV=0', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb'),
    'new_baseline.kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.kpp.lmd.005': (
        'KPP ν0=2.5, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods'
    )
}
datasets = {
    "TAO": tao_gridded,
    # "MITgcm": mitgcm,
}
%autoreload

with warnings.catch_warnings():
    warnings.simplefilter("ignore", category=UserWarning)
    warnings.simplefilter("ignore", category=FutureWarning)
    for short_name, (long_name, folder) in tqdm.tqdm(catalog_sub.items()):
        datasets.update(
            {
                short_name: mixpods.load_mom6_sections(folder).assign_attrs(
                    {"title": long_name}
                )
            }
        )
100%|██████████| 5/5 [00:31<00:00,  6.39s/it]
datasets["les"] = les["0.0.-140.oct.average.month"].to_dataset()

# Offset LES to work with slicing below
datasets["les"]["time"] = datasets["les"]["time"] + pd.Timedelta(days=25 * 365)
tree = DataTree.from_dict(datasets)
tree
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*

Modify catalog selection?#

if "les" in tree:
    tree["les"] = tree["les"].isel(z=slice(-2))
    tree["les"]["KT"].attrs["standard_name"] = "ocean_vertical_heat_diffusivity"

if "micro" in locals():
    tree.update(micro)

Post-process catalog subset#

tree = tree.sel(time=slice("2000", "2017"))
tree
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*
ref = tree["TAO"].ds.reset_coords(drop=True).cf.sel(Z=slice(-120, None))
counts = np.minimum(ref["S2"].cf.count("Z"), ref["N2T"].cf.count("Z")).load()


def calc_histograms(ds):
    ds = ds.copy()
    ds["tao_mask"] = counts.reindex(time=ds.time, method="nearest") > 5
    ds["tao_mask"].attrs = {
        "description": "True when there are more than 5 5-m T, u, v in TAO dataset"
    }
    # ds = ds.where(ds.tao_mask)
    return ds.update(mixpods.pdf_N2S2(ds))


tree = tree.map_over_subtree(calc_histograms)

daily composite#

%autoreload

dailies = tree.map_over_subtree(mixpods.daily_composites)
dailies
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/numpy/lib/nanfunctions.py:1217: RuntimeWarning: All-NaN slice encountered
  r, k = function_base._ureduce(a, func=_nanmedian, axis=axis, out=out,
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*

seasonal cycle#

tree = mixpods.persist_tree(tree)

Validate before continuing#

mixpods.validate_tree(tree)

TODO: Verify depth is normalized#

for name, ds in datasets.items():
    (ds.cf["sea_water_x_velocity"].cf["Z"].plot(marker=".", ls="none", label=name))
plt.legend()
<matplotlib.legend.Legend>
_images/2372712bb196933da3d938a29d3f04e0e217fe29d04bb036dbceb08ab9ce12da.png

Mean profiles#

%autoreload

S2 = mixpods.plot_profile_fill(tree, "S2", "S^2")
N2 = mixpods.plot_profile_fill(tree, "N2T", "N_T^2")
u = mixpods.plot_profile_fill(tree, "sea_water_x_velocity", "u")
T = mixpods.plot_profile_fill(tree, "sea_water_potential_temperature", "T")
Ri = mixpods.plot_median_Ri(tree)
(S2 + N2 + Ri + u + T).cols(3)

Stability Diagram#

%autoreload

mixpods.plot_stability_diagram_by_dataset(tree, nrows=2)
_images/67078def4b6e6abed5d4d3b9f4c9ced685ef78f679212ea748097dffb73b5b30.png

Daily composites#

%autoreload

mixpods.plot_daily_composites(dailies, ["eps", "KT"], logy=True, legend=False).opts(
    hv.opts.GridSpace(show_legend=True)
)
%autoreload

mixpods.plot_daily_composites(
    dailies, ["S2", "N2", "Rig_T"], logy=False, legend=False
).opts(hv.opts.GridSpace(show_legend=True, width=200), hv.opts.Overlay(frame_width=200))

Turbulence distributions#

handles = [
    mixpods.plot_distributions(tree, "chi", bins=np.linspace(-11, -4, 101), log=True),
    mixpods.plot_distributions(tree, "eps", bins=np.linspace(-11, -4, 101), log=True),
    mixpods.plot_distributions(
        tree, "ocean_vertical_heat_diffusivity", bins=np.linspace(-8, -1, 101), log=True
    ),
    # plot_distributions(tree, "Jq", bins=np.linspace(-1000, 0, 51), log=False),
    mixpods.plot_distributions(tree, "Rig_T", np.linspace(-0.5, 1.5, 61))
    * hv.VLine(0.25).opts(line_color="k"),
]
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/computation.py:769: RuntimeWarning: invalid value encountered in log10
  result_data = func(*input_data)
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/computation.py:769: RuntimeWarning: invalid value encountered in log10
  result_data = func(*input_data)
hv.Layout(handles).opts("Overlay", frame_width=600).cols(2)

Compare boundary layer depth#

mixing_layer_depth_criteria = {
    "boundary_layer_depth": {"name": "KPPhbl|KPP_OBLdepth|ePBL_h_ML"},
}

hbl = (
    tree.drop_nodes("TAO")
    .dc.subset_nodes("KPP_OBLdepth")
    .dc.concatenate_nodes()
    .reset_coords(drop=True)
).load()
(
    # hbl.groupby("time.hour").mean().hvplot.line(by="node", flip_yaxis=True)
    hbl.groupby("time.hour").median().hvplot.line(by="node", flip_yaxis=True)
    + hbl.to_dataframe().hvplot.hist(
        by="node",
        bins=np.arange(0, 90, 1),
        normed=1,
        alpha=0.3,
        ylim=(0, 0.05),
        muted_alpha=0,
    )
).opts(hv.opts.Curve(invert_yaxis=True))

Heat Budget#

f, ax = plt.subplots(
    2,
    math.ceil(len(tree) / 2),
    constrained_layout=True,
    squeeze=False,
    sharex=True,
    sharey=True,
    figsize=(10, 3),
)

for axx, (name, node) in zip(ax.flat, tree.children.items()):
    mixpods.plot_climo_heat_budget_1d(node.ds, mxldepth=-40, penetration="mom", ax=axx)
    axx.set_title(name)

dcpy.plots.clean_axes(ax)